import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader, dataloader
import pandas as pd
from models.autoencoder import autoencoder
from utils import *
import pandas as pd
import matplotlib.pyplot as plt
import math as math

T = 62.8
def X_Y_new():
    # 测试数据
    raw_data = np.load('data.npy')
    miss_index = np.linspace(0, raw_data.shape[0] - 1, 100).astype('int')

    miss_data = raw_data[miss_index]

    x1 = np.cos(2 * math.pi * miss_index/T).reshape(-1,1)
    x2 = np.sin(2 * math.pi * miss_index/T).reshape(-1,1)
    miss_index = np.concatenate((x1,x2),axis=1)

    miss_data = np.expand_dims(miss_data, 1)
    return miss_index,miss_data

def X_Y():
    # 测试数据
    raw_data = np.load('data.npy')
    miss_index = np.linspace(0, raw_data.shape[0] - 1, 100).astype('int')

    miss_data = raw_data[miss_index]
    miss_index = normalization(miss_index)


    miss_index = np.expand_dims(miss_index, 1)
    miss_data = np.expand_dims(miss_data, 1)
    return miss_index,miss_data



 

if __name__ == '__main__':

    X,Y = X_Y_new()   # 添加周期性的XY
    # X,Y = X_Y()    # 普通的XY


    AE = autoencoder(X.shape[1],Y.shape[1])
    criterion = nn.MSELoss(size_average=False)
    optimizer = torch.optim.Adam(AE.parameters(), lr=1e-3)

    for e in range(2000):
        img = torch.from_numpy(X).to(torch.float32)
        label = torch.from_numpy(Y).to(torch.float32)
        optimizer.zero_grad()
        # 前向传播
        _, output = AE(img)
        loss = criterion(output, label)
        # 反向传播
        loss.backward()
        optimizer.step()
        if (e+1) % 50 == 0: # 每 20 次，输出loss
            print('epoch: {}, Loss: {:.4f}'.format(e + 1, loss.data))

    # # 测试网络
    newtimepoints_scaled = np.expand_dims(np.linspace(0, 300-1, 300), 1)
    # choose new:
    x1 = np.cos(2 * math.pi * newtimepoints_scaled/T).reshape(-1,1)
    x2 = np.sin(2 * math.pi * newtimepoints_scaled/T).reshape(-1,1)
    newtimepoints_scaled = np.concatenate((x1,x2),axis=1)


    # else:
    # newtimepoints_scaled = normalization(newtimepoints_scaled)


    _, output = AE(torch.from_numpy(newtimepoints_scaled).to(torch.float32))

    np.save('r.npy', output.detach().numpy())

